"""
Qwen3-4B-AWQ模型API测试工具

此模块提供了针对已部署的Qwen3-4B-AWQ模型的测试方法，包括API调用、错误处理和测试用例。
模型服务器地址: 192.168.1.236:8000
"""

import requests
import json
import time
from typing import Dict, List, Optional, Union, Any
import logging
from dataclasses import dataclass


# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


@dataclass
class ModelConfig:
    """模型配置类"""
    host: str = "192.168.1.236"
    port: int = 8000
    model_name: str = ""  # 将在运行时获取实际模型名称
    api_endpoint: str = "/v1/chat/completions"
    models_endpoint: str = "/v1/models"  # 模型列表端点
    health_endpoint: str = "/health"  # 健康检查端点
    timeout: int = 30  # 请求超时时间（秒）
    max_retries: int = 3  # 最大重试次数
    retry_delay: float = 1.0  # 重试延迟（秒）
    
    # 生成参数默认值（优化后的值）
    max_tokens: int = 200  # 减少默认最大token数以提高速度
    temperature: float = 0.1  # 降低默认温度以减少随机性，提高速度
    top_p: float = 0.9  # 核采样参数
    frequency_penalty: float = 0.0  # 频率惩罚
    presence_penalty: float = 0.0  # 存在惩罚
    
    # 流式输出配置
    stream: bool = False  # 是否使用流式输出
    
    # 请求优化配置
    use_cache: bool = True  # 是否使用缓存
    skip_special_tokens: bool = True  # 是否跳过特殊token


class QwenModelClient:
    """Qwen3-4B-AWQ模型API客户端"""
    
    def __init__(self, config: Optional[ModelConfig] = None):
        """
        初始化模型客户端
        
        Args:
            config: 模型配置，如果为None则使用默认配置
        """
        self.config = config or ModelConfig()
        self.base_url = f"http://{self.config.host}:{self.config.port}"
        self.api_url = f"{self.base_url}{self.config.api_endpoint}"
        self.models_url = f"{self.base_url}{self.config.models_endpoint}"
        self.health_url = f"{self.base_url}{self.config.health_endpoint}"
        
        # 设置默认请求头
        self.headers = {
            "Content-Type": "application/json",
            "Accept": "application/json"
        }
        
        # 初始化时获取可用模型列表
        self._available_models = []
        self._refresh_models()
    
    def _prepare_request_data(self, messages: List[Dict[str, str]], 
                             **kwargs) -> Dict[str, Any]:
        """
        准备请求数据
        
        Args:
            messages: 对话消息列表
            **kwargs: 其他参数
            
        Returns:
            格式化的请求数据字典
        """
        request_data = {
            "model": self.config.model_name,
            "messages": messages,
            "max_tokens": kwargs.get("max_tokens", self.config.max_tokens),
            "temperature": kwargs.get("temperature", self.config.temperature),
            "stream": kwargs.get("stream", False)
        }
        
        # 添加其他可选参数
        optional_params = ["top_p", "frequency_penalty", "presence_penalty", "stop"]
        for param in optional_params:
            if param in kwargs:
                request_data[param] = kwargs[param]
                
        return request_data
    
    def _handle_response(self, response: requests.Response) -> Dict[str, Any]:
        """
        处理API响应
        
        Args:
            response: HTTP响应对象
            
        Returns:
            解析后的响应数据
            
        Raises:
            ValueError: 当响应格式不正确时
            requests.HTTPError: 当服务器返回错误状态码时
        """
        try:
            response.raise_for_status()
            return response.json()
        except requests.HTTPError as http_err:
            logger.error(f"HTTP错误: {http_err}")
            logger.error(f"响应状态码: {response.status_code}")
            logger.error(f"响应内容: {response.text}")
            raise
        except json.JSONDecodeError as json_err:
            logger.error(f"JSON解析错误: {json_err}")
            logger.error(f"响应内容: {response.text}")
            raise ValueError("服务器返回了无效的JSON格式")
        except Exception as err:
            logger.error(f"处理响应时发生未知错误: {err}")
            raise
    
    def chat_completion(self, messages: List[Dict[str, str]], 
                       **kwargs) -> Dict[str, Any]:
        """
        调用模型进行对话补全
        
        Args:
            messages: 对话消息列表，格式为[{"role": "user", "content": "你好"}]
            **kwargs: 其他生成参数
            
        Returns:
            模型生成的响应
            
        Raises:
            requests.exceptions.ConnectionError: 网络连接错误
            requests.exceptions.Timeout: 请求超时
            requests.HTTPError: HTTP错误
            ValueError: 响应格式错误
        """
        try:
            request_data = self._prepare_request_data(messages, **kwargs)
            logger.info(f"发送请求到: {self.api_url}")
            logger.debug(f"请求数据: {json.dumps(request_data, ensure_ascii=False)}")
            
            response = requests.post(
                self.api_url,
                headers=self.headers,
                json=request_data,
                timeout=self.config.timeout
            )
            
            return self._handle_response(response)
            
        except requests.exceptions.ConnectionError as conn_err:
            logger.error(f"连接错误: 无法连接到服务器 {self.config.host}:{self.config.port}")
            logger.error(f"详细错误: {conn_err}")
            raise
        except requests.exceptions.Timeout as timeout_err:
            logger.error(f"请求超时: 服务器在 {self.config.timeout} 秒内未响应")
            logger.error(f"详细错误: {timeout_err}")
            raise
        except requests.exceptions.RequestException as req_err:
            logger.error(f"请求异常: {req_err}")
            raise
        except Exception as err:
            logger.error(f"未知错误: {err}")
            raise
    
    def fast_chat(self, prompt: str, max_tokens: int = 100) -> Dict[str, Any]:
        """
        快速聊天模式，使用优化参数提高响应速度
        
        Args:
            prompt: 用户输入的提示
            max_tokens: 最大生成token数（默认100）
            
        Returns:
            模型响应
        """
        messages = [{"role": "user", "content": prompt}]
        
        # 使用优化的参数
        return self.chat_completion(
            messages=messages,
            max_tokens=max_tokens,
            temperature=0.1,  # 低温度减少随机性，提高速度
            top_p=0.8,  # 稍微降低top_p值
            use_cache=True,  # 启用缓存
            skip_special_tokens=True  # 跳过特殊token
        )
    
    def simple_chat(self, prompt: str, fast_mode: bool = True, **kwargs) -> str:
        """
        简化的单轮对话接口
        
        Args:
            prompt: 用户输入的提示
            fast_mode: 是否使用快速模式（默认True）
            **kwargs: 其他生成参数
            
        Returns:
            模型生成的文本回复
        """
        if fast_mode:
            # 使用快速模式，限制token数量
            max_tokens = kwargs.get('max_tokens', 100)
            response = self.fast_chat(prompt, max_tokens=max_tokens)
        else:
            # 使用标准模式
            messages = [{"role": "user", "content": prompt}]
            response = self.chat_completion(messages, **kwargs)
        
        try:
            return response["choices"][0]["message"]["content"]
        except (KeyError, IndexError, TypeError) as e:
            logger.error(f"解析响应内容失败: {e}")
            logger.error(f"响应数据: {response}")
            raise ValueError("无法从响应中提取生成的内容")
    
    def _refresh_models(self):
        """
        刷新可用模型列表
        """
        try:
            response = requests.get(self.models_url, timeout=self.config.timeout)
            if response.status_code == 200:
                models_data = response.json()
                self._available_models = [model["id"] for model in models_data.get("data", [])]
                logger.info(f"已获取模型列表: {self._available_models}")
                
                # 如果配置的模型名称不在可用列表中，使用第一个可用模型
                if self.config.model_name not in self._available_models and self._available_models:
                    self.config.model_name = self._available_models[0]
                    logger.info(f"使用模型: {self.config.model_name}")
            else:
                logger.warning(f"获取模型列表失败，状态码: {response.status_code}")
        except Exception as e:
            logger.warning(f"获取模型列表异常: {e}")
    
    def get_available_models(self) -> List[str]:
        """
        获取可用模型列表
        
        Returns:
            模型ID列表
        """
        return self._available_models.copy()
    
    def health_check(self) -> bool:
        """
        检查模型服务是否健康
        
        Returns:
            如果服务正常返回True，否则返回False
        """
        try:
            # 使用/health端点进行健康检查
            response = requests.get(self.health_url, timeout=5)
            if response.status_code == 200:
                logger.info("健康检查通过 (/health端点)")
                return True
        except Exception as e:
            logger.debug(f"/health端点检查失败: {e}")
        
        try:
            # 备用方案：使用/ping端点进行健康检查
            ping_url = f"{self.base_url}/ping"
            response = requests.get(ping_url, timeout=5)
            if response.status_code == 200:
                logger.info("健康检查通过 (/ping端点)")
                return True
        except Exception as e:
            logger.debug(f"/ping端点检查失败: {e}")
        
        try:
            # 最后备用方案：发送一个简单的测试请求
            test_messages = [{"role": "user", "content": "Hi"}]
            self.chat_completion(test_messages, max_tokens=10)
            logger.info("健康检查通过 (测试请求)")
            return True
        except Exception as e:
            logger.error(f"所有健康检查均失败: {e}")
            return False


if __name__ == "__main__":
    # 简单测试示例
    client = QwenModelClient()
    
    try:
        # 健康检查
        if client.health_check():
            print("模型服务健康检查通过")
            
            # 简单对话测试
            response = client.simple_chat("你好，请介绍一下你自己")
            print(f"模型回复: {response}")
        else:
            print("模型服务健康检查失败")
    except Exception as e:
        print(f"测试过程中发生错误: {e}")